import chromadb
from chromadb.config import Settings
from embedding_client import EmbeddingClient
import uuid

class VectorStore:
    def __init__(self,collection_name,embedding_client : EmbeddingClient):
        #内存模式
        chroma_client = chromadb.Client(Settings(allow_reset=True))
        #数据持久化
        #chroma_client = chromadb.PersistentClient(path="./chroma")
        #演示使用，实际不要用，清空了
        chroma_client.reset()
        #创建一个connection
        self.collection = chroma_client.get_or_create_collection(name=collection_name)
        self.embedding_client = embedding_client

    def add_documents(self,documents):
        embeddings = self.embedding_client.get_embeddings(documents)
        #print("documents:",documents,"embeddings:",embeddings)
        '''向collection中添加文档与向量'''
        self.collection.add(
            embeddings = embeddings,  # 每个文档的向量
            documents = documents , #文档的原文
            ids=[str(uuid.uuid4()) for _ in range(len(documents))]
        )

    def search(self, query, top_n):
        '''检索向量数据库'''
        results = self.collection.query(
            query_embeddings = self.embedding_client.get_embeddings([query]),
            n_results =  top_n
        )
        return results

